{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center>\n",
    "    <p style=\"text-align:center\">\n",
    "        <img alt=\"phoenix logo\" src=\"https://storage.googleapis.com/arize-phoenix-assets/assets/phoenix-logo-light.svg\" width=\"200\"/>\n",
    "        <br>\n",
    "        <a href=\"https://docs.arize.com/phoenix/\">Docs</a>\n",
    "        |\n",
    "        <a href=\"https://github.com/Arize-ai/phoenix\">GitHub</a>\n",
    "        |\n",
    "        <a href=\"https://join.slack.com/t/arize-ai/shared_invite/zt-1px8dcmlf-fmThhDFD_V_48oU7ALan4Q\">Community</a>\n",
    "    </p>\n",
    "</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Instrumenting AWS Bedrock client with OpenInference and Phoenix\n",
    "\n",
    "In this tutorial we will trace model calls to AWS Bedrock using OpenInference. The OpenInference Bedrock tracer instruments the Python `boto3` library, so all `invoke_model` calls will automatically generate traces that can be sent to Phoenix.\n",
    "\n",
    "ℹ️ This notebook requires a valid AWS configuration and access to AWS Bedrock and the `claude-v2` model from Anthropic."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Install dependencies and set up OpenTelemetry tracer\n",
    "\n",
    "First install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install arize-phoenix boto3 openinference-instrumentation-bedrock"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from urllib.parse import urljoin\n",
    "\n",
    "import boto3\n",
    "import phoenix as px\n",
    "from openinference.instrumentation.bedrock import BedrockInstrumentor\n",
    "from opentelemetry import trace as trace_api\n",
    "from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter\n",
    "from opentelemetry.sdk import trace as trace_sdk\n",
    "from opentelemetry.sdk.resources import Resource\n",
    "from opentelemetry.sdk.trace.export import ConsoleSpanExporter, SimpleSpanProcessor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Start a Pheonix server to collect traces. Be sure to view Phoenix in your browser to watch traces show up in Phoenix as they are collected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🌍 To view the Phoenix app in your browser, visit http://localhost:6006/\n",
      "📺 To view the Phoenix app in a notebook, run `px.active_session().view()`\n",
      "📖 For more information on how to use Phoenix, check out https://docs.arize.com/phoenix\n"
     ]
    }
   ],
   "source": [
    "px.launch_app()\n",
    "session_url = px.active_session().url"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we're configuring the OpenTelemetry tracer by adding two SpanProcessors. The first SpanProcessor will simply print all traces received from OpenInference instrumentation to the console. The second will export traces to Phoenix so they can be collected and viewed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "resource = Resource(attributes={})\n",
    "tracer_provider = trace_sdk.TracerProvider(resource=resource)\n",
    "span_console_exporter = ConsoleSpanExporter()\n",
    "phoenix_otlp_endpoint = urljoin(session_url, \"v1/traces\")\n",
    "phoenix_exporter = OTLPSpanExporter(endpoint=phoenix_otlp_endpoint)\n",
    "tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter=span_console_exporter))\n",
    "tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter=phoenix_exporter))\n",
    "trace_api.set_tracer_provider(tracer_provider=tracer_provider)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Instrumenting Bedrock clients\n",
    "\n",
    "Now, let's create a `boto3` session. This initiates a configured environment for interacting with AWS services. If you haven't yet configured `boto3` to use your credentials, please refer to the [official documentation](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html). Or, if you have the AWS CLI, run `aws configure` from your terminal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "session = boto3.session.Session()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Clients created using this session configuration are currently uninstrumented. We'll make one for comparison."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "uninstrumented_client = session.client(\"bedrock-runtime\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we instrument Bedrock with our OpenInference instrumentor. All Bedrock clients created after this call will automatically produce traces when calling `invoke_model`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "BedrockInstrumentor().instrument()\n",
    "instrumented_client = session.client(\"bedrock-runtime\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Calling the LLM and viewing OpenInference traces\n",
    "\n",
    "Calling `invoke_model` using the `uninstrumented_client` will produce no traces, but will show the output from the LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " I'm doing well, thanks for asking!\n"
     ]
    }
   ],
   "source": [
    "prompt = b'{\"prompt\": \"Human: Hello there, how are you? Assistant:\", \"max_tokens_to_sample\": 1024}'\n",
    "response = uninstrumented_client.invoke_model(modelId=\"anthropic.claude-v2\", body=prompt)\n",
    "response_body = json.loads(response.get(\"body\").read())\n",
    "print(response_body[\"completion\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LLM calls using the `instrumented_client` will print traces to the console! By configuring the `SpanProcessor` to export to a different OpenTelemetry collector, your OpenInference spans can be collected and analyzed to better understand the behavior of your LLM application."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "    \"name\": \"bedrock.invoke_model\",\n",
      "    \"context\": {\n",
      "        \"trace_id\": \"0x994497f3948acc4cf390658d6f43dd7c\",\n",
      "        \"span_id\": \"0x3180fe54cacc8f9b\",\n",
      "        \"trace_state\": \"[]\"\n",
      "    },\n",
      "    \"kind\": \"SpanKind.INTERNAL\",\n",
      "    \"parent_id\": null,\n",
      "    \"start_time\": \"2024-02-15T15:52:13.188122Z\",\n",
      "    \"end_time\": \"2024-02-15T15:52:14.664294Z\",\n",
      "    \"status\": {\n",
      "        \"status_code\": \"UNSET\"\n",
      "    },\n",
      "    \"attributes\": {\n",
      "        \"llm.prompts\": \"Human: Hello there, how are you? Assistant:\",\n",
      "        \"llm.invocation_parameters\": \"{\\\"max_tokens_to_sample\\\": 1024}\",\n",
      "        \"llm.model_name\": \"anthropic.claude-v2\",\n",
      "        \"message.content\": \" I'm doing well, thanks for asking!\"\n",
      "    },\n",
      "    \"events\": [],\n",
      "    \"links\": [],\n",
      "    \"resource\": {\n",
      "        \"attributes\": {},\n",
      "        \"schema_url\": \"\"\n",
      "    }\n",
      "}\n",
      " I'm doing well, thanks for asking!\n"
     ]
    }
   ],
   "source": [
    "response = instrumented_client.invoke_model(modelId=\"anthropic.claude-v2\", body=prompt)\n",
    "response_body = json.loads(response.get(\"body\").read())\n",
    "print(response_body[\"completion\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "More information about our instrumentation integrations, OpenInerence can be found in our [documentation](https://docs.arize.com/phoenix/telemetry/instrumentation)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dev38",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
